﻿#include "estree_visitor.hpp"
#include "dummy_node.h"

#include <charconv>

namespace cyclone::parser
{

    static auto formatCallChainAncestor(ParserRuleContext *node, const Json &ancestor)
    {
        Json CallChain;
        JsonVisitorHelper::formatManualNode(CallChain, node, KeyValue{"type", "ChainExpression"});
        CallChain["expression"] = ancestor;
        return CallChain;
    }
    static auto isCallChainElement(antlr4::tree::ParseTree *ct)
    {
        return dynamic_cast<FormalParser::CallExpressionContext *>(ct) ||
               dynamic_cast<FormalParser::MemberExpressionContext *>(ct);
    };
    template <typename T>
    static auto checkChainElementIsAncestor(T *node)
    {
        FormalParser::ExpressionContext *parent = node;
        while (parent)
        {
            if (!parent->parent || !isCallChainElement(parent->parent))
                break;
            parent = static_cast<FormalParser::ExpressionContext *>(parent->parent);
        }
        return parent;
    }

    bool EstreeVisitor::isOptionalCallChain(FormalParser::ExpressionContext *ctx)
    {
        return optionalCallChainAncestors_.find(ctx) != optionalCallChainAncestors_.end();
    }

    Json EstreeVisitor::visit(FormalParser::MemberExpressionContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        auto jsonifyNode = JsonVisitorHelper::jsonifyRuleNode(ctx);
        if (ctx->object_)
            jsonifyNode["object"] = dispatch(ctx->object_);
        if (ctx->property_)
            jsonifyNode["property"] = dispatch(ctx->property_);
        jsonifyNode["computed"] = ctx->property_->computed;
        jsonifyNode["optional"] = ctx->property_->optional;
        if (ctx->property_->optional)
            optionalCallChainAncestors_.insert(checkChainElementIsAncestor(ctx));
        return isOptionalCallChain(ctx) ? formatCallChainAncestor(ctx, jsonifyNode) : jsonifyNode;
    }
    Json EstreeVisitor::visit(FormalParser::UnaryExpressionContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        auto jsonifyNode = JsonVisitorHelper::jsonifyRuleNode(ctx);
        if (ctx->operator_)
            jsonifyNode["prefix"] = ctx->operator_->prefix;
        if (ctx->operator_->children.size())
            jsonifyNode["operator"] = reinterpret_cast<TerminalNode *>(ctx->operator_->children[0])->getText();
        jsonifyNode["argument"] = dispatch(ctx->argument_);
        return jsonifyNode;
    }
    Json EstreeVisitor::visit(FormalParser::UpdateExpressionContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        auto jsonifyNode = JsonVisitorHelper::jsonifyRuleNode(ctx);
        if (ctx->operator_)
            jsonifyNode["prefix"] = ctx->operator_->prefix;
        if (ctx->operator_->children.size())
            jsonifyNode["operator"] = reinterpret_cast<TerminalNode *>(ctx->operator_->children[0])->getText();
        if (ctx->argument_)
            jsonifyNode["argument"] = dispatch(ctx->argument_);
        return jsonifyNode;
    }
    Json EstreeVisitor::visit(FormalParser::SourceElementContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        IF_NULLPTR_RETURN_DEFAULT(ctx->statement_);
        return dispatch(ctx->statement_);
    }
    Json EstreeVisitor::visit(FormalParser::ProgramContext *ctx)
    {
        return JsonVisitor::visit(ctx);
    }

    Json EstreeVisitor::visit(FormalParser::DotMemberExpressionContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        return visitAnonymousNode(ctx);
    }

    Json EstreeVisitor::visit(FormalParser::BracketMemberExpressionContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        return visitAnonymousNode(ctx->expression_);
    }

    Json EstreeVisitor::visit(FormalParser::IdentifierExpressionContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        return visitAnonymousNode(ctx);
    }

    Json EstreeVisitor::visit(FormalParser::FunctionExpressionContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        auto jsonifyNode = JsonVisitorHelper::jsonifyRuleNode(ctx);
        if (ctx->baseFunction_)
        {
            jsonifyNode["async"] = ctx->baseFunction_->async_;
            jsonifyNode["generator"] = ctx->baseFunction_->generator_;
            jsonifyNode["params"] = ctx->baseFunction_->params_ ? dispatch(ctx->baseFunction_->params_) : JsonArray();
            if (ctx->baseFunction_->id_)
                jsonifyNode["id"] = dispatch(ctx->baseFunction_->id_);
        }
        jsonifyNode["body"] = JsonVisitor::visit(ctx->body_);
        jsonifyNode["expression"] = false;
        return jsonifyNode;
    }

    Json EstreeVisitor::visit(FormalParser::YieldExpressionContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        auto jsonifyNode = JsonVisitorHelper::jsonifyRuleNode(ctx);
        jsonifyNode["delegate"] = nullptr != ctx->Multiply();
        jsonifyNode["argument"] = dispatch(ctx->argument_);
        return jsonifyNode;
    }

    Json EstreeVisitor::visit(FormalParser::PrivateMemberExpressionContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        return visit(ctx->privateIdentifier());
    }

    Json EstreeVisitor::visit(FormalParser::PrivateIdentifierContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        auto jsonifyNode = JsonVisitorHelper::jsonifyRuleNode(ctx);
        jsonifyNode["name"] = ctx->name;
        return jsonifyNode;
    }
    Json EstreeVisitor::visit(FormalParser::ExpressionSequenceContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        if (ctx->expression().size() == 1)
            return dispatch(ctx->expression()[0]);
        auto jsonifyNode = JsonVisitorHelper::jsonifyRuleNode(ctx);
        auto array = Json::array();
        for (const auto &expr : ctx->expression())
            array.push_back(dispatch(expr));
        jsonifyNode["expressions"] = array;
        return jsonifyNode;
    }

    Json EstreeVisitor::visit(FormalParser::IdentifierContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        auto jsonifyNode = JsonVisitorHelper::jsonifyRuleNode(ctx);
        JsonVisitor::visit(ctx);
        if (ctx->Identifier())
            jsonifyNode["name"] = ctx->Identifier()->getText();
        return processTypeAnnotationInPattern(ctx, jsonifyNode);
    }
    Json EstreeVisitor::visit(FormalParser::ObjectPatternContext *ctx)
    {
        return processTypeAnnotationInPattern(ctx);
    }
    Json EstreeVisitor::visit(FormalParser::ArrayPatternContext *ctx)
    {
        return processTypeAnnotationInPattern(ctx);
    }
    Json EstreeVisitor::visit(FormalParser::RestElementContext *ctx)
    {
        return processTypeAnnotationInPattern(ctx);
    }

    Json EstreeVisitor::visit(FormalParser::CallExpressionContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        auto jsonifyNode = JsonVisitorHelper::jsonifyRuleNode(ctx);
        jsonifyNode["callee"] = dispatch(ctx->callee_);
        jsonifyNode["arguments"] = dispatch(ctx->arguments_);
        auto isOptional = nullptr != ctx->OptionalChainingOperator();
        jsonifyNode["optional"] = isOptional;
        if (isOptional)
            optionalCallChainAncestors_.insert(checkChainElementIsAncestor(ctx));
        return isOptionalCallChain(ctx) ? formatCallChainAncestor(ctx, jsonifyNode) : jsonifyNode;
    }

    Json EstreeVisitor::visit(FormalParser::LiteralExpressionContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        return visit(ctx->literal());
    }

    Json EstreeVisitor::visit(FormalParser::LiteralContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        // NOTE check if literal is templateString
        // if (auto templateLiteral = ctx->templateLiteral())
        //     return visit(templateLiteral);
        auto jsonifyNode = JsonVisitorHelper::jsonifyRuleNode(ctx);
        jsonifyNode["type"] = "Literal";
        // NOTE need to remove '' or " " for sting
        if (auto stringLiteral = ctx->StringLiteral())
            jsonifyNode["value"] = JsonVisitorHelper::removeQuoteFromString(ctx->getText());
        else if (auto number = ctx->numericLiteral())
        {
            if (auto decimal = number->DecimalLiteral())
            {
                double value;
                std::from_chars(&ctx->getText().at(0), &ctx->getText().at(ctx->getText().size() - 1), value);
                jsonifyNode["value"] = value;
            }
            else
            {
                long long value;
                std::from_chars(&ctx->getText().at(0), &ctx->getText().at(ctx->getText().size() - 1), value);
                jsonifyNode["value"] = value;
            }
        }
        // null
        else if (auto literal = ctx->NullLiteral())
        {
            jsonifyNode["value"] = Json();
        }
        // true or false
        else if (auto literal = ctx->BooleanLiteral())
        {
            if (("true" == literal->getText()))
                jsonifyNode["value"] = true;
            else
                jsonifyNode["value"] = false;
        }
        else
            jsonifyNode["value"] = ctx->getText();
        jsonifyNode["raw"] = ctx->getText();
        return jsonifyNode;
    }

    Json EstreeVisitor::visit(FormalParser::ThisExpressionContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        return JsonVisitorHelper::jsonifyRuleNode(ctx);
    }

    Json EstreeVisitor::visit(FormalParser::ArrowFunctionExpressionContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        auto jsonifyNode = JsonVisitorHelper::jsonifyRuleNode(ctx);
        jsonifyNode["async"] = nullptr != ctx->Async();
        // Cannot be generator itself
        jsonifyNode["generator"] = false;
        // Identifier must be null
        jsonifyNode["id"] = Json();
        jsonifyNode["params"] = ctx->params_ ? dispatch(ctx->params_) : JsonArray();
        jsonifyNode["body"] = dispatch(ctx->body_);
        jsonifyNode["expression"] = ctx->body_ && ctx->body_->expression();
        return jsonifyNode;
    }

    Json EstreeVisitor::visit(FormalParser::TemplateLiteralContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        auto jsonifyNode = JsonVisitorHelper::jsonifyRuleNode(ctx);
        JsonArray quasis;
        JsonArray expressions;
        for (const auto &atom : ctx->templateStringAtom())
            if (is<FormalParser::BraceTemplateStringContext>(atom))
                expressions.push_back(dispatch(reinterpret_cast<FormalParser::BraceTemplateStringContext *>(atom)->expression_));
        const auto allQuasis = TemplateElement::formatTemplateElement(ctx);
        JsonArray quasisCollection;
        for (const auto &quasis : allQuasis)
        {
            auto jsonifyQuasis = JsonVisitorHelper::jsonifyRuleNode(*quasis);
            jsonifyQuasis["tail"] = quasis->tail;
            Json quasisValue;
            quasisValue["cooked"] = quasis->value.cooked.has_value() ? quasis->value.cooked.value() : "";
            quasisValue["raw"] = quasis->value.raw;
            jsonifyQuasis["value"] = quasisValue;
            quasisCollection.push_back(jsonifyQuasis);
        };
        jsonifyNode["quasis"] = quasisCollection;
        jsonifyNode["expressions"] = expressions;
        return jsonifyNode;
    }

    Json EstreeVisitor::visit(FormalParser::ClassExpressionContext *ctx)
    {
        // TODO
        return JsonVisitor::visit(ctx);
    }

    Json EstreeVisitor::visit(FormalParser::PropertyShorthandContext *ctx)
    {
        auto key = JsonVisitor::visit(ctx);
        key["value"] = key.find("key") != key.end() ? key["key"] : Json();
        return key;
    }
    Json EstreeVisitor::visit(FormalParser::ShorthandAssignmentPropertyContext *ctx)
    {
        auto key = JsonVisitor::visit(ctx);
        key["value"] = key.find("key") != key.end() ? key["key"] : Json();
        return key;
    }

    Json EstreeVisitor::visit(FormalParser::VariableDeclaratorContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        auto jsonifyNode = JsonVisitorHelper::jsonifyRuleNode(ctx);
        jsonifyNode["id"] = dispatch(ctx->id_);
        // if (ctx->annotation_)
        //     jsonifyNode["annotation"] = dispatch(ctx->annotation_);
        jsonifyNode["init"] = dispatch(ctx->init_);
        return jsonifyNode;
    }

    Json EstreeVisitor::visit(FormalParser::FunctionDeclarationContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        auto jsonifyNode = JsonVisitorHelper::jsonifyRuleNode(ctx);
        if (ctx->baseFunction_)
        {
            jsonifyNode["async"] = ctx->baseFunction_->async_;
            jsonifyNode["generator"] = ctx->baseFunction_->generator_;
            if (ctx->baseFunction_->id_)
                jsonifyNode["id"] = dispatch(ctx->baseFunction_->id_);
            jsonifyNode["params"] = ctx->baseFunction_->params_ ? dispatch(ctx->baseFunction_->params_) : JsonArray();
        }
        jsonifyNode["body"] = JsonVisitor::visit(ctx->body_);
        jsonifyNode["expression"] = false;
        return jsonifyNode;
    }

    Json EstreeVisitor::visit(FormalParser::ReturnStatementContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        auto jsonifyNode = JsonVisitorHelper::jsonifyRuleNode(ctx);
        jsonifyNode["argument"] = dispatch(ctx->argument_);
        return jsonifyNode;
    }

    Json EstreeVisitor::visit(FormalParser::ParenthesizedExpressionContext *ctx)
    {
        return this->preserveParens() ? JsonVisitor::visit(ctx) : dispatch(ctx->expression_);
    }

    Json EstreeVisitor::visit(FormalParser::MetaPropertyIdentifierContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        auto jsonifyNode = JsonVisitorHelper::jsonifyRuleNode(ctx);
        jsonifyNode["name"] = ctx->name;
        return jsonifyNode;
    }
    Json EstreeVisitor::visit(FormalParser::LogicalExpressionContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        auto jsonifyNode = JsonVisitor::visit(ctx);
        auto op = ctx->DoubleQuestionMark() ? ctx->DoubleQuestionMark() : (ctx->And() ? ctx->And() : ctx->Or());
        jsonifyNode["operator"] = op ? op->getText() : "";
        return jsonifyNode;
    }

    Json EstreeVisitor::visit(DummyNode *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        auto jsonifyNode = JsonVisitorHelper::jsonifyRuleNode(ctx);
        jsonifyNode["name"] = ctx->toString();
        return jsonifyNode;
    }

    Json EstreeVisitor::visit(FormalParser::PrimitiveTypeContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        return JsonVisitorHelper::jsonifyRuleNode(ctx);
    }

    Json EstreeVisitor::visit(FormalParser::TupleTypeContext *ctx)
    {
        IF_NULLPTR_RETURN_DEFAULT(ctx);
        auto jsonifyNode = JsonVisitorHelper::jsonifyRuleNode(ctx);
        auto elements = Json::array();
        if (ctx->elements)
            for (auto e : ctx->elements->type_())
                elements.push_back(dispatch(e));
        return jsonifyNode;
    }

}
